import re
def split_long_lines(lines, max_length=60):
    new_lines = []
    for line in lines:
        if len(line) <= max_length:
            new_lines.append(line)
        else:
            # 按标点符号切分长行
            parts = re.split(r'([，。；！？])', line)
            current_line = ""
            for i in range(0, len(parts), 2):
                if i + 1 < len(parts):
                    part = parts[i] + parts[i + 1]
                else:
                    part = parts[i]
                if len(current_line) + len(part) <= max_length:
                    current_line += part
                else:
                    if current_line:
                        new_lines.append(current_line)
                    current_line = part
            if current_line:
                new_lines.append(current_line)
    return new_lines


def merge_short_lines(lines, min_length=20):
    new_lines = []
    i = 0
    while i < len(lines):
        current_line = lines[i]
        while i + 1 < len(lines) and len(current_line) < min_length:
            current_line += " " + lines[i + 1]
            i += 1
        new_lines.append(current_line)
        i += 1
    return new_lines


def preprocess_lyrics(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.read().splitlines()
    # 切分长行
    lines = split_long_lines(lines)
    # 合并短行
    lines = merge_short_lines(lines)
    return lines


if __name__ == "__main__":
    file_path = "lyrics.txt"
    processed_lines = preprocess_lyrics(file_path)
    output_file = "train.txt"

    with open(output_file, 'w', encoding='utf-8') as file:
        for line in processed_lines:
            file.write(line + '\n')
    print(f"处理后的歌词已成功写入 {output_file}。")



